Skip to content

[Common] Split grouped activation build - #3430

Open
harryzhou2000 wants to merge 1 commit into
NVIDIA:mainfrom
harryzhou2000:hhanyu/grouped-activation-tu-split
Open

[Common] Split grouped activation build#3430
harryzhou2000 wants to merge 1 commit into
NVIDIA:mainfrom
harryzhou2000:hhanyu/grouped-activation-tu-split

Conversation

@harryzhou2000

Copy link
Copy Markdown
Contributor

Description

Split five monolithic grouped-activation CUDA translation units into smaller
per-activation and per-direction translation units:

  • GELU and QGELU forward, backward, and dbias
  • ReLU and SReLU forward, backward, and dbias
  • SiLU and SwiGLU forward and backward

The existing function bodies and exported symbols are moved without changing
kernel logic or public APIs. swiglu_grouped_dbias.cu remains unchanged because
it contains only one instantiation.

Build performance

The comparison used separate empty ccache directories with zero hits on the same
full-node GCP-NRT allocation (224 logical CPUs), container, venv, source base,
and submodules. MAX_JOBS and CMAKE_BUILD_PARALLEL_LEVEL were unset, so Ninja
used machine parallelism. Both builds used NVTE_BUILD_THREADS_PER_JOB=4,
NVTE_CUDA_ARCHS=100, and the PyTorch-only configuration. The split build ran
first, so any system cache warming favored the baseline.

Metric Baseline Split Reduction Speedup
Maximum translation unit 932.579 s 445.884 s 52.2% 2.09x
Common-library build 935.905 s 448.261 s 52.1% 2.09x
Full editable install 1,038.390 s 551.548 s 46.9% 1.88x

The split increases aggregate compiler work by 13.7% because template parsing
is repeated across more translation units. With machine parallelism, that work
moves off the critical path: the longest translation unit changes from
gelu_grouped.cu to the unrelated fused_topk_with_score_function.cu.

Validation

  • Focused pre-commit checks: merge-conflict, large-file, EOF, trailing-whitespace,
    and clang-format checks passed
  • git diff --check passed
  • Confirmed all 14 moved functions remain exported
  • Expanded grouped MXFP8 activation test matrix: 915 passed, 810 expected skips,
    0 failures

Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
@greptile-apps

greptile-apps Bot commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR splits grouped activation CUDA implementations into smaller per-activation and per-direction translation units without changing their APIs or kernel dispatch logic.

  • Registers all new translation units in the architecture-specific CMake source list.
  • Preserves optional fast-math compilation for every moved implementation.
  • Retains each grouped activation and dbias function exactly once with unchanged declarations and exports.

Confidence Score: 5/5

The PR appears safe to merge, with the moved APIs, build registration, export surface, and source-specific compilation settings preserved.

The change only redistributes existing grouped activation definitions, and all moved functions remain uniquely defined and included in the same architecture and fast-math build paths.

Important Files Changed

Filename Overview
transformer_engine/common/CMakeLists.txt Registers all split CUDA translation units in the architecture-specific and optional fast-math lists, preserving their prior compilation properties.
transformer_engine/common/activation/gelu_grouped.cu Retains GELU forward while moving the unchanged GELU backward and QGELU implementations into dedicated translation units.
transformer_engine/common/activation/qgelu_grouped.cu Defines the relocated QGELU forward API with unchanged dispatch behavior.
transformer_engine/common/activation/relu_grouped.cu Retains ReLU forward while moving unchanged ReLU backward and SReLU implementations into dedicated translation units.
transformer_engine/common/activation/srelu_grouped.cu Defines the relocated SReLU forward API with unchanged dispatch behavior.
transformer_engine/common/activation/swiglu_grouped_bwd.cu Defines the relocated grouped SiLU backward API with unchanged arguments and template dispatch.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[Grouped activation APIs] --> B[Per-activation CUDA translation units]
    B --> C[Forward implementations]
    B --> D[Backward implementations]
    B --> E[Dbias implementations]
    C --> F[Architecture-specific CMake source list]
    D --> F
    E --> F
    F --> G[Transformer Engine shared library]
Loading

Reviews (1): Last reviewed commit: "[Common] Split grouped activation build" | Re-trigger Greptile

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant